module ub_wr_arbiter (
    input logic clk,
    input logic rst_n,
    gdma_ub_wr_req_if.in u_arb_normal_wr_req_if,
    gdma_ub_wr_req_if.in u_arb_atomic_wr_req_if,
    gdma_ub_wr_req_if.in u_gdma_wr_req_if,
    gdma_ub_wr_req_if.out u_gnt_wr_req_if  // Output uses GDMA format (128B)
);

    // Priority: Atomic (highest) > Round-robin between ARU normal and GDMA
    // Round-robin state for normal and GDMA requests
    logic            last_grant_gdma;  // 0: last was ARU normal, 1: last was GDMA

    // ARU atomic request (128B, direct pass-through)
    logic            aru_atomic_req_vld;
    gdma_ub_wr_req_s aru_atomic_req_pld;

    always_comb begin
        aru_atomic_req_vld = u_arb_atomic_wr_req_if.vld;
        aru_atomic_req_pld = u_arb_atomic_wr_req_if.pld;
    end

    // ARU normal request (128B, direct pass-through)
    logic            aru_normal_req_vld;
    gdma_ub_wr_req_s aru_normal_req_pld;

    always_comb begin
        aru_normal_req_vld = u_arb_normal_wr_req_if.vld;
        aru_normal_req_pld = u_arb_normal_wr_req_if.pld;
    end

    // GDMA request (128B, direct pass-through)
    logic            gdma_req_vld;
    gdma_ub_wr_req_s gdma_req_pld;

    always_comb begin
        gdma_req_vld = u_gdma_wr_req_if.vld;
        gdma_req_pld = u_gdma_wr_req_if.pld;
    end

    // Arbitration logic
    logic [1:0] grant_sel;  // 0: ARU normal, 1: GDMA, 2: ARU atomic
    logic       grant_vld;

    always_comb begin
        // Priority 1: Atomic ALU (highest priority, never backpressure)
        if (aru_atomic_req_vld) begin
            grant_sel = 2'd2;
            grant_vld = 1'b1;
        end  // Priority 2: Round-robin between ARU normal and GDMA
        else if (last_grant_gdma) begin
            // Last was GDMA, try ARU normal first
            if (aru_normal_req_vld) begin
                grant_sel = 2'd0;
                grant_vld = 1'b1;
            end else if (gdma_req_vld) begin
                grant_sel = 2'd1;
                grant_vld = 1'b1;
            end else begin
                grant_sel = 2'd0;
                grant_vld = 1'b0;
            end
        end else begin
            // Last was ARU normal (or initial), try GDMA first
            if (gdma_req_vld) begin
                grant_sel = 2'd1;
                grant_vld = 1'b1;
            end else if (aru_normal_req_vld) begin
                grant_sel = 2'd0;
                grant_vld = 1'b1;
            end else begin
                grant_sel = 2'd0;
                grant_vld = 1'b0;
            end
        end
    end

    // Update round-robin state
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            last_grant_gdma <= 1'b0;
        end else begin
            // Update round-robin state
            if (grant_vld && u_gnt_wr_req_if.rdy) begin
                if (grant_sel == 2'd0) begin
                    // ARU normal completed
                    last_grant_gdma <= 1'b0;
                end else if (grant_sel == 2'd1) begin
                    // GDMA completed
                    last_grant_gdma <= 1'b1;
                end
                // ARU atomic (grant_sel==2'd2) doesn't update round-robin
            end
        end
    end

    // Output assignment
    always_comb begin
        u_gnt_wr_req_if.vld = grant_vld;
        case (grant_sel)
            2'd0: u_gnt_wr_req_if.pld = aru_normal_req_pld;
            2'd1: u_gnt_wr_req_if.pld = gdma_req_pld;
            2'd2: u_gnt_wr_req_if.pld = aru_atomic_req_pld;
            default: u_gnt_wr_req_if.pld = '0;
        endcase
    end

    // Backpressure ready signals
    // ARU normal: 128B request, assert rdy immediately
    assign u_arb_normal_wr_req_if.rdy = (grant_sel == 2'd0) && grant_vld && u_gnt_wr_req_if.rdy;

    // GDMA: 128B request, assert rdy immediately
    assign u_gdma_wr_req_if.rdy       = (grant_sel == 2'd1) && grant_vld && u_gnt_wr_req_if.rdy;

    // ARU atomic: 128B request, assert rdy immediately
    assign u_arb_atomic_wr_req_if.rdy = (grant_sel == 2'd2) && grant_vld && u_gnt_wr_req_if.rdy;

endmodule
